__global__
void kernel_gemm_impl1(float const * const __restrict lhs, float const * const __restrict rhs, float * const __restrict out,
                 ::std::uint32_t const M, ::std::uint32_t const K, ::std::uint32_t const N,
                 bool const is_lhs_transpose,
                 bool const is_rhs_transpose,
                 bool const is_out_transpose)
{
    ::std::uint32_t const thread_idx{blockIdx.x * blockDim.x + threadIdx.x};
#if 1
    __builtin_assume(thread_idx < N);
#else
    if (__builtin_expect(thread_idx >= N, 0)) {
        __builtin_abort();
    }
#endif

    tops_dte_ctx_t dte_lhs;
    tops_dte_ctx_t dte_rhs;

    // 初始化DTE
    // FIXME use dte_scope instead
    dte_lhs.init();
    dte_rhs.init();

    for (::std::uint32_t n{thread_idx}; n < N; n += /* thread_num = */ 24)
    {
        __valigned__ float l1_rhs[20'000];
        ::tops::mdspan l1_rhs_view(tops::Private, l1_rhs, K);

        if (is_rhs_transpose) {
            ::tops::mdspan l3_rhs_view(tops::Global, const_cast<float*>(rhs + n * K), K);
            ::tops::memcpy(dte_rhs, l1_rhs_view, l3_rhs_view);
        } else {
            for (::std::uint32_t k{}; k < K; ++k) {
                ::tops::mdspan l3_input_r_view(tops::Global, const_cast<float*>(rhs + k * N + n), 1);
                ::tops::mdspan l1_rhs_indexing_k_view(tops::Private, l1_rhs + k, 1);
                ::tops::memcpy(dte_rhs, l1_rhs_indexing_k_view, l3_input_r_view);
            }
        }

        for (::std::uint32_t m{}; m < M; ++m)
        {
            float output{};

            __valigned__ float l1_lhs[20'000];
            ::tops::mdspan l1_lhs_view(tops::Private, l1_lhs, K);

            if (is_lhs_transpose) {
                for (auto k = 0; k < K; ++k) {
                    ::tops::mdspan l3_input_l_view(tops::Global, const_cast<float*>(lhs + k * M + m), 1);
                    ::tops::mdspan l1_lhs_indexing_k_view(tops::Private, l1_lhs + k, 1);
                    ::tops::memcpy(dte_lhs, l1_lhs_indexing_k_view, l3_input_l_view);
                }
            } else {
                ::tops::mdspan l3_lhs_view(tops::Global, const_cast<float*>(lhs + m * K), K);
                ::tops::memcpy(dte_lhs, l1_lhs_view, l3_lhs_view);
            }

            ::std::uint32_t index{};
#if 1
            auto v_result = ::tops::vbroadcast<vfloat>(0.f);
            for (; index + 32 <= K; index += /* tops::vlength<vfloat>() = */ 32)
            {
                auto const v_lhs = ::tops::vload<vfloat>(l1_lhs + index);
                auto const v_rhs = ::tops::vload<vfloat>(l1_rhs + index);
                v_result += v_lhs * v_rhs;
            }

            float list_result[32];
            ::tops::vstore(v_result, list_result);
            for (::std::uint32_t i{}; i < 32; ++i) {
                output += list_result[i];
            }
#endif
            for (;index < K; ++index) {
                output += l1_lhs[index] * l1_rhs[index];
            }

            ::tops::mdspan l1_output_view(tops::Private, __builtin_addressof(output), 1);
            if (is_out_transpose) {
                ::tops::mdspan l3_output_view(tops::Global, out + n * M + m, 1);
                ::tops::memcpy(dte_rhs, l3_output_view, l1_output_view);
            } else {
                ::tops::mdspan l3_output_view(tops::Global, out + m * N + n, 1);
                ::tops::memcpy(dte_rhs, l3_output_view, l1_output_view);
            }
        }
    }
}

__global__
void kernel_gemm_impl2(float const * const __restrict lhs, float const * const __restrict rhs, float * const __restrict out,
                 ::std::uint32_t const M, ::std::uint32_t const K, ::std::uint32_t const N,
                 bool const is_lhs_transpose,
                 bool const is_rhs_transpose,
                 bool const is_out_transpose)
{
    ::std::uint32_t const thread_idx{blockIdx.x * blockDim.x + threadIdx.x};

    tops_dte_ctx_t dte;

    // 初始化DTE
    // FIXME use dte_scope instead
    dte.init();

    ::std::uint32_t const loop_end{M * N};
    for (::std::uint32_t i{thread_idx}; i < loop_end; i += 24)
    {
        ::std::uint32_t const m {i / N};
        ::std::uint32_t const n {i % N};

        float output{};

        for (auto k = 0; k < K; ++k)
        {
            float l1_input_l [[clang::uninitialized]];
            ::tops::mdspan l1_input_l_view(tops::Private, __builtin_addressof(l1_input_l), 1);
            if (is_lhs_transpose) {
                ::tops::mdspan l3_input_l_view(tops::Global, const_cast<float*>(lhs + k * M + m), 1);
                ::tops::memcpy(dte, l1_input_l_view, l3_input_l_view);
            } else {
                ::tops::mdspan l3_input_l_view(tops::Global, const_cast<float*>(lhs + m * K + k), 1);
                ::tops::memcpy(dte, l1_input_l_view, l3_input_l_view);
            }

            float l1_input_r [[clang::uninitialized]];
            ::tops::mdspan l1_input_r_view(tops::Private, __builtin_addressof(l1_input_r), 1);
            if (is_rhs_transpose) {
                ::tops::mdspan l3_input_r_view(tops::Global, const_cast<float*>(rhs + n * K + k), 1);
                ::tops::memcpy(dte, l1_input_r_view, l3_input_r_view);
            } else {
                ::tops::mdspan l3_input_r_view(tops::Global, const_cast<float*>(rhs + k * N + n), 1);
                ::tops::memcpy(dte, l1_input_r_view, l3_input_r_view);
            }

            output += l1_input_l * l1_input_r;
        }

        ::tops::mdspan l1_output_view(tops::Private, __builtin_addressof(output), 1);
        if (is_out_transpose) {
            ::tops::mdspan l3_output_view(tops::Global, out + n * M + m, 1);
            ::tops::memcpy(dte, l3_output_view, l1_output_view);
        } else {
            ::tops::mdspan l3_output_view(tops::Global, out + m * N + n, 1);
            ::tops::memcpy(dte, l3_output_view, l1_output_view);
        }
    }
}

void GCU_GEMM(float * __restrict dev_lhs,
              float * __restrict dev_rhs,
              float * __restrict dev_out,
              int m, int k, int n,
              bool is_lhs_transpose,
              bool is_rhs_transpose,
              bool is_out_transpose)
{
    if (n <= 24) {
        ::kernel_gemm_impl2<<<2, 12>>>(dev_lhs, dev_rhs, dev_out,
            ::std::uint32_t(m), ::std::uint32_t(k), ::std::uint32_t(n),
            is_lhs_transpose,
            is_rhs_transpose,
            is_out_transpose);
    } else {
        ::kernel_gemm_impl1<<<2, 12>>>(dev_lhs, dev_rhs, dev_out,
            ::std::uint32_t(m), ::std::uint32_t(k), ::std::uint32_t(n),
                            is_lhs_transpose,
                            is_rhs_transpose,
                            is_out_transpose);
    }
}
